package main

import (
	"fmt"
	"GA_ASCON/ascon_helper"
	"GA_ASCON/hdf5_helper"
	"GA_ASCON/ga_helper"
	"math/rand"
	"time"
	"os"
	"strconv"
)

func main(){
	h5path := os.Args[1]
	keys,nonce0ss,nonce1ss:=hdf5_helper.Readh5(h5path)
	path := os.Args[2]
	// predict := hdf5_helper.Readpredict_r(path) //回归模型
	predict := hdf5_helper.Readpredict_c(path,2) //2分类模型
	// predict := hdf5_helper.Readpredict_c(path,65)  //65分类模型
	// flag_2  := false
	n_traces,_ := strconv.Atoi(os.Args[3])
	fmt.Printf("Key:%v\n",keys[0])
	
	//初始化密钥种群
	n_pop :=16
	var r_key [][]uint8
	for i:=0;i<n_pop;i++{
		tmp:=make([]uint8,8)
		for j:=0;j<8;j++{
			tmp[j]=uint8(rand.Intn(256))
		}
		r_key = append(r_key,tmp)
	}
	
	//初始化fitntess
	fit_key := make([]ga_helper.Fit_Key,n_pop)

	for ge:=0;ge<200;ge++{
		//随机选择
		index_traces := make([]int,10000)
		for i:=0;i<len(index_traces);i++{
			index_traces[i]=i
		}
		s := rand.New(rand.NewSource(time.Now().UnixNano()))
		s.Shuffle(len(index_traces),func(i,j int){index_traces[i],index_traces[j]=index_traces[j],index_traces[i]})
		index_traces = index_traces[:n_traces]
		nonce0s,nonce1s := nonce0ss[50000:60000],nonce1ss[50000:60000]
		for i:=0;i<len(index_traces);i++{
			nonce0s[i] = nonce0s[index_traces[i]]
			nonce1s[i] = nonce1s[index_traces[i]]
		}
		nonce0s = nonce0s[0:n_traces]
		nonce1s = nonce1s[0:n_traces]

		predicts := predict
		for i:=0;i<len(index_traces);i++{
			predict[i] = predict[index_traces[i]]
		}
		predicts = predicts[:n_traces]
		//开始遗传算法
		ga_helper.Cross(r_key,0.5)
		// ga_helper.Mutation(r_key,0.05,nonce0s,nonce1s,predicts,ga_helper.Coef) //回归模型
		ga_helper.Mutation_c(r_key,0.05,nonce0s,nonce1s,predicts,ga_helper.Score,true) //2-分类模型
		// ga_helper.Mutation_c(r_key,0.05,nonce0s,nonce1s,predicts,ga_helper.Score,false) //65-分类模型
		for i:=0;i<n_pop;i++{
			labels_y := ascon_helper.Computer_Sbox_labels(r_key[i],nonce0s,nonce1s,64)
			Hws := ascon_helper.HW_2c(labels_y)         //2分类模型
			// Hws := ascon_helper.Hw_weights(labels_y) //65分类和回归
			fit_key[i] = ga_helper.Fit_Key{ga_helper.Score(Hws,predicts),r_key[i]} //分类模型 
			// fit_key[i] = ga_helper.Fit_Key{ga_helper.Coef(Hws,predicts),r_key[i]} //回归模型 
		}
		
		ga_helper.Selection(r_key,fit_key)
		Result := ga_helper.Result(keys[0],fit_key)[0]
		fmt.Printf("第%d代适应度排序:%v\n",ge,Result)
		if Result.Num_bytes == 8{
			fmt.Printf("第%d代后找到正确密钥：%v\n",ge,Result.Key)
			return 
		}
	}
}